from Dataloader.twitterloader import *
from SentModel.Sent2Vec import *
from SentModel.sentence_trainer import *
from torch.utils.data import DataLoader
import pickle

tr = LMReader("../data/twitter_LM_tr.txt")
dev = LMReader("../data/twitter_LM_dev.txt")
te = LMReader("../data/twitter_LM_te.txt")

batchsize_GPU = 16
GPU_Count = torch.cuda.device_count()
train_batchsize = batchsize_GPU*GPU_Count
tr_loader = DataLoader(tr, batch_size=train_batchsize, shuffle=True)
dev_loader = DataLoader(dev, batch_size=100, shuffle=True)
te_loader = DataLoader(te, batch_size=100, shuffle=True)

with open("../saved/config.pkl", "rb") as fr:
    config = pickle.load(fr)
model = LSTMVec("../saved/bert_embedding.pkl", "../../bert_en/", 768, 768, 2)
trainer = LMTrainer(model, config)
trainer.LMTrain(tr_loader, dev_loader, te_loader, max_epoch=3, print_every=10, learning_rate=2e-3, model_file="../saved/lstm_lm.pkl")
print("============Finetune With Embedding Parameters======================")
trainer.sent2vec.emb_update = True
trainer.LMTrain(tr_loader, dev_loader, te_loader, max_epoch=5, print_every=10, learning_rate=2e-5, model_file="../saved/lstm_lm.pkl")